from .exceptions import IndexErrorOfStructure


class UF1:
    """使用数组实现并查集"""
    """
    @is_connected  O(1)
    @union_elements  O(n)
    """
    def __init__(self, length):
        self.ID = list(range(length))

    def size(self):
        return len(self.ID)

    def __len__(self):
        return len(self.ID)

    def _find(self, p):
        """查看元素p的集合编号"""
        if p < 0 or p >= len(self.ID):
            raise IndexErrorOfStructure("Index out of range.")
        return self.ID[p]

    def is_connected(self, p, q):
        """查看元素p和q是否同属一个集合"""
        return self._find(p) == self._find(q)

    def union_elements(self, p, q):
        """合并p, q所在的集合"""
        p_id = self._find(p)
        q_id = self._find(q)
        if p_id == q_id:
            return
        i = 0
        while i < len(self.ID):
            if self.ID[i] == p_id:
                self.ID[i] = q_id
            i += 1


class UF2:
    def __init__(self, length):
        self.parent = list(range(length))

    def __len__(self):
        return len(self.parent)

    def size(self):
        return len(self.parent)

    def _find(self, p):
        """返回一个节点的根节点的索引"""
        if p < 0 or p >= len(self.parent):
            raise IndexErrorOfStructure("Index out of range.")
        while p != self.parent[p]:
            p = self.parent[p]
        return p

    def is_connected(self, p, q):
        """查看元素p和q是否同属一个集合"""
        return self._find(p) == self._find(q)

    def union_elements(self, p, q):
        """合并p, q所在的集合"""
        p_root = self._find(p)
        q_root = self._find(q)
        if p_root == q_root:
            return
        self.parent[p_root] = q_root


class UF3:
    def __init__(self, length):
        self.parent = list(range(length))
        # num中存储的是以对应索引位置树的节点总数
        self.num = list(1 for _ in range(length))

    def __len__(self):
        return len(self.parent)

    def size(self):
        return len(self.parent)

    def _find(self, p):
        """返回一个节点的根节点的索引"""
        if p < 0 or p >= len(self.parent):
            raise IndexErrorOfStructure("Index out of range.")
        while p != self.parent[p]:
            p = self.parent[p]
        return p

    def is_connected(self, p, q):
        """查看元素p和q是否同属一个集合"""
        return self._find(p) == self._find(q)

    def union_elements(self, p, q):
        """合并p, q所在的集合"""
        p_root = self._find(p)
        q_root = self._find(q)
        if p_root == q_root:
            return
        if self.num[p_root] < self.num[q_root]:
            self.parent[p_root] = q_root
            self.num[q_root] += self.num[p_root]
        else:
            self.parent[q_root] = p_root
            self.num[q_root] += self.num[p_root]


class UF4:
    def __init__(self, length):
        self.parent = list(range(length))
        # rank中存储的是以对应索引位置树的高度
        self.rank = list(1 for _ in range(length))

    def __len__(self):
        return len(self.parent)

    def size(self):
        return len(self.parent)

    def _find(self, p):
        """返回一个节点的根节点的索引"""
        if p < 0 or p >= len(self.parent):
            raise IndexErrorOfStructure("Index out of range.")
        while p != self.parent[p]:
            p = self.parent[p]
        return p

    def is_connected(self, p, q):
        """查看元素p和q是否同属一个集合"""
        return self._find(p) == self._find(q)

    def union_elements(self, p, q):
        """合并p, q所在的集合"""
        p_root = self._find(p)
        q_root = self._find(q)
        if p_root == q_root:
            return
        if self.rank[p_root] < self.rank[q_root]:
            self.parent[p_root] = q_root
        elif self.rank[p_root] > self.rank[q_root]:
            self.parent[q_root] = p_root
        else:
            self.parent[q_root] = p_root
            self.rank[q_root] += 1


class UF5:
    def __init__(self, length):
        self.parent = list(range(length))
        # rank中存储的是以对应索引位置树的高度
        self.rank = list(1 for _ in range(length))

    def __len__(self):
        return len(self.parent)

    def size(self):
        return len(self.parent)

    def _find(self, p):
        """返回一个节点的根节点的索引"""
        if p < 0 or p >= len(self.parent):
            raise IndexErrorOfStructure("Index out of range.")
        while p != self.parent[p]:
            # 进行路径压缩，让当前节点指向他父亲节点的父亲节点
            # 注意添加了压缩操作后，rank不在是高度，但是仍然可以使用，相当于排名
            self.parent[p] = self.parent[self.parent[p]]
            p = self.parent[p]
        return p

    def is_connected(self, p, q):
        """查看元素p和q是否同属一个集合"""
        return self._find(p) == self._find(q)

    def union_elements(self, p, q):
        """合并p, q所在的集合"""
        p_root = self._find(p)
        q_root = self._find(q)
        if p_root == q_root:
            return
        if self.rank[p_root] < self.rank[q_root]:
            self.parent[p_root] = q_root
        elif self.rank[p_root] > self.rank[q_root]:
            self.parent[q_root] = p_root
        else:
            self.parent[q_root] = p_root
            self.rank[q_root] += 1
